#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   logger.py
@Contact :   xxzhang16@fudan.edu.cn

@Modify Time      @Author    @Version    @Desciption
------------      -------    --------    -----------
2022/3/28 14:01   zxx      1.0         None
'''

# import lib


class Logger:
    def __init__(self, is_finetune=False, log_step=50):
        """
        log_step: TrainLoop evaluate in every epoch, but loss log need to specify
            default: conv-20, rec-50

        Useful values in the whole trainning procedure are most included below.
        loss：
            train rec:
                informax_pretrain (IP):
                    IP_info_db_loss
                    IP_joint_loss = info_db_loss
                recomendation:
                    rec_loss
                    info_db_loss
                    joint_loss = rec_loss+0.025*info_db_loss
            train gen:
                gen_loss
                selection_loss
                joint_loss = gen_loss_weight * gen_loss + selection_loss
        metrics：
            gen:
                ppl
                dist1
                dist2
                dist3
                dist4
                bleu1
                bleu2
                bleu3
                bleu4
                count
                ReR@1
                ReR@10
                ReR@50
                true_recall_movie_count,
                res_movie_recall
            rec:
                ReR@1
                ReR@10
                ReR@50
                loss
                gate
                count
                gate_count
        """
        self.is_finetune = is_finetune
        self.log_step = log_step

        if not is_finetune:
            self.log_dict = {
                # distinguish which part of the model was trained
                'module': 'info and rec',

                # log step
                'log_step': self.log_step,

                # loss part
                'loss': {
                    'IP_info_db_loss': [],
                    'IP_joint_loss': [],
                    'rec_loss': [],
                    'info_db_loss': [],
                    'joint_loss': []
                },

                # metric part,
                'metric': {
                    'recall@1': [],
                    'recall@10': [],
                    'recall@50': [],
                    'loss': [],
                    'gate': [],
                    'count': [],
                    'gate_count': []
                },

                'best_metric': {
                    'recall@1': -1,
                    'recall@10': -1,
                    'recall@50': -1,
                    'loss': -1,
                    'gate': -1,
                    'count': -1,
                    'gate_count': -1
                }
            }
        else:
            self.log_dict = {
                # distinguish which part of the model was trained
                'module': 'gen',

                # log step
                'log_step': self.log_step,

                # loss part
                'loss': {
                    'gen_loss': [],
                    'selection_loss': [],
                    'joint_loss': []
                },

                # metric part,
                'metric': {
                    'ppl': [],
                    'dist1': [],
                    'dist2': [],
                    'dist3': [],
                    'dist4': [],
                    'bleu1': [],
                    'bleu2': [],
                    'bleu3': [],
                    'bleu4': [],
                    'count': [],
                    'true_recall_movie_count': [],
                    'res_movie_recall': [],
                    'recall@1': [],
                    'recall@10': [],
                    'recall@50': []
                },
                'best_metric': {
                    'ppl': -1,
                    'dist1': -1,
                    'dist2': -1,
                    'dist3': -1,
                    'dist4': -1,
                    'bleu1': -1,
                    'bleu2': -1,
                    'bleu3': -1,
                    'bleu4': -1,
                    'count': -1,
                    'true_recall_movie_count': -1,
                    'res_movie_recall': -1,
                    'recall@1': -1,
                    'recall@10': -1,
                    'recall@50': -1
                }
            }

    def log_loss(self, loss_name, loss_value):
        self.log_dict['loss'][loss_name].append(loss_value)

    def log_metric(self, metric_dict):
        for key in metric_dict.keys():
            self.log_dict['metric'][key].append(metric_dict[key])

    def log_best_metric(self, best_metric_dict):
        for key in best_metric_dict.keys():
            self.log_dict['best_metric'][key] = best_metric_dict[key]

    def write_to_file(self, log_save_pth):
        import json
        with open(log_save_pth, 'w') as file_writer:
            json.dump(self.log_dict, file_writer)


if __name__ == '__main__':
    # unit test
    # rec
    import torch
    rec_logger = Logger(is_finetune=False, log_step=50)
    for i in range(300):
        if i % 50 == 0:
            rec_logger.log_loss('IP_info_db_loss', i / 300)
            rec_logger.log_loss('IP_joint_loss', i / 300)

    for i in range(1000):
        if i % 50 == 0:
            rec_logger.log_loss('info_db_loss', i / 1000)
            rec_logger.log_loss('rec_loss', i / 1000)
            rec_logger.log_loss('joint_loss', i / 1000)
        if i % 250 == 0:
            metric = {
                'recall@1': i / 5,
                'recall@10': i / 25,
                'recall@50': i / 50,
                'loss': i / 1000,
                'gate': i / 100,
                'count': i,
                'gate_count': i / 500
            }
            # metric = {key: torch.tensor(metric[key]) for key in metric.keys()}

            rec_logger.log_metric(metric)
            if i == 500:
                rec_logger.log_best_metric(metric)
    rec_logger.write_to_file('tmp_rec_log_file.json')

    # gen
    # gen_logger = Logger(is_finetune=True, log_step=20)
    # for i in range(1000):
    #     if i % 20 == 0:
    #         gen_logger.log_loss('gen_loss', i / 1000)
    #         gen_logger.log_loss('selection_loss', i / 1000)
    #         gen_logger.log_loss('joint_loss', i / 1000)
    #     if i % 100 == 0:
    #         metric = {
    #             "ppl": i / 10,
    #             "dist1": i / 20,
    #             "dist2": i / 30,
    #             "dist3": i / 40,
    #             "dist4": i / 50,
    #             "bleu1": i / 60,
    #             "bleu2": i / 70,
    #             "bleu3": i / 80,
    #             "bleu4": i / 90,
    #             "count": i / 100,
    #             "true_recall_movie_count": i / 110,
    #             "res_movie_recall": i / 120,
    #             "recall@1": i / 1000,
    #             "recall@10": i / 500,
    #             "recall@50": i / 100}
    #         gen_logger.log_metric(metric)
    #         if i == 500:
    #             gen_logger.log_best_metric(metric)
    # gen_logger.write_to_file('tmp_gen_log_file.json')
